import cv2
import os
import subprocess
import imagehash
from PIL import Image
from tqdm import tqdm

def extract_i_frames(input_path, output_dir):
    """使用FFmpeg精确提取I帧"""
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    cmd = [
        'ffmpeg',
        '-i', input_path,
        '-vf', "select='eq(pict_type,I)'",
        '-vsync', 'vfr',
        '-q:v', '2',
        os.path.join(output_dir, 'iframe_%04d.jpg')
    ]
    subprocess.run(cmd, check=True)

def remove_similar_frames(frame_dir, threshold=5):
    """基于感知哈希去除相似帧"""
    hashes = {}
    duplicates = 0
    
    for frame in tqdm(sorted(os.listdir(frame_dir))):
        img_path = os.path.join(frame_dir, frame)
        try:
            with Image.open(img_path) as img:
                h = imagehash.dhash(img)
                match_found = False
                for existing_hash in hashes.keys():
                    if h - existing_hash < threshold:
                        match_found = True
                        break
                if not match_found:
                    hashes[h] = img_path
                else:
                    os.remove(img_path)
                    duplicates += 1
        except Exception as e:
            print(f"处理 {frame} 时出错: {e}")
    
    print(f"移除 {duplicates} 张相似帧，保留 {len(hashes)} 张关键帧")

def scene_change_detection(input_path, output_dir, threshold=0.3):
    """基于场景变化的帧提取"""
    cap = cv2.VideoCapture(input_path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    prev_frame = None
    frame_count = 0
    
    while True:
        ret, frame = cap.read()
        if not ret:
            break
            
        if frame_count % int(fps) == 0 or prev_frame is None:  # 每秒至少取一帧
            cv2.imwrite(f"{output_dir}/scene_{frame_count:06d}.jpg", frame)
            prev_frame = frame
            frame_count += 1
            continue
            
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        prev_gray = cv2.cvtColor(prev_frame, cv2.COLOR_BGR2GRAY)
        diff = cv2.absdiff(gray, prev_gray)
        diff_percent = diff.sum() / (gray.shape[0] * gray.shape[1] * 255)
        
        if diff_percent > threshold:
            cv2.imwrite(f"{output_dir}/scene_{frame_count:06d}.jpg", frame)
            prev_frame = frame
            frame_count += 1
    
    cap.release()

def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('input', help='输入视频路径')
    parser.add_argument('-o', '--output', default='keyframes', help='输出目录')
    parser.add_argument('-t', '--threshold', type=int, default=5, 
                       help='哈希相似度阈值(0-64，值越小越严格)')
    args = parser.parse_args()
    
    # 第一阶段：提取所有I帧
    temp_dir = os.path.join(args.output, "temp_frames")
    extract_i_frames(args.input, temp_dir)
    
    # 第二阶段：去重处理
    remove_similar_frames(temp_dir, args.threshold)
    
    # 重命名最终文件
    final_dir = os.path.join(args.output, "final_frames")
    os.makedirs(final_dir, exist_ok=True)
    for i, frame in enumerate(sorted(os.listdir(temp_dir))):
        os.rename(
            os.path.join(temp_dir, frame),
            os.path.join(final_dir, f"keyframe_{i:04d}.jpg")
        )
    os.rmdir(temp_dir)
    print(f"关键帧已保存至: {final_dir}")

if __name__ == '__main__':
    main()
